import torch
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from Config import config   #从自定义的Config.py参数文件中输入
from Model import Model
from sklearn.manifold import TSNE

#加载参数
config = config()

# ToTensor():1、转化为一个tensor；2、将数值归一化，转换到0-1之间；3、将channel维度放在第一维度上
transformation = transforms.Compose([
    transforms.ToTensor()
])

test_ds = datasets.MNIST(config.data_path, train=False, download=True, transform=transformation)
test_dl = DataLoader(test_ds, batch_size=config.batch_size, shuffle=True)

#绘制训练前的tsne图
test_data = test_ds.data.view(-1, 28*28).numpy()    #将图像展平为向量
#使用t-SNE进行降维
tsne = TSNE(n_components=2, random_state=42)
embeded_data = tsne.fit_transform(test_data)
#绘制模型训练前的 t-SNE 图
plt.figure(figsize=(10, 8))
# embeded_data[:, 0], 取降维后的数据的第一维(轴)
# embeded_data[:, 1], 取降维后的数据的第二维(y轴)
# c=test_ds.targets, 使用test_ds.target作为颜色编码，test_ds.target包含每个样本的真实标签，不同的标签会用不同的颜色表示
# cmap='viridis' 指定颜色的映射方案viridis是一种连续的颜色映射，也可以选择其他的颜色映射：plasma, inferno, maga等
plt.scatter(embeded_data[:, 0], embeded_data[:, 1], c=test_ds.targets, cmap='viridis')
plt.colorbar()
plt.title('t-SNE Visualization before Model Training')
plt.xlabel('t-SNE Dimension 1')
plt.ylabel('t-SNE Dimension 2')
plt.show()

# 绘制训练后的tsne图
# 载入模型和参数
model = Model()
model.load_state_dict(torch.load(config.save_path))     #导入网格的参数
#设置为评估模型，保留所有的神经元节点
# dropout层：在训练时，dropout层会随机丢弃一部分神经元以防止过拟合，在评估模式下，所有神经元都会被保留，输出是确定性的
# batch normalization层：在训练时，batchnorm层使用当前batch的统计信息(均值和方差)进行归一化，
#                       而在评估模式下，它会使用整个训练过程中的累积的全局统计信息
model = model.eval()    #消除dropout层的影响

#初始化一个空列表features用于存储输出的特征向量，每个特征向量是模型对输出图像的高级表示
features = []
# 存储每个向量对应的真实标签
labels = []
#将图片输入到模型中，输出预测矩阵
#禁用梯度(进入一个上下文管理器)，在这个中所有的张量操作都不会记录梯度信息，为了节省内存
with torch.no_grad():
    # images 是一个形状为[batch_size,channel,height,width]的张量，表示图像
    # true 是形状为 [batch_size]的张量，表示这批图像对应的真实标签
    for images, true_labels in test_dl:
        outputs = model(images)     #将图像的模型获取并输出
        features.append(outputs)
        labels.append(true_labels)
# torch.cat函数作为拼接存在
# 将features列表中的所有张量沿着第0维(即批次维度)拼接成一个大的张量
# dim=0 表示沿着批次的维度进行拼接
features = torch.cat(features, dim=0)
labels = torch.cat(labels, dim=0)
# 将特征向量转换为 numpy数组
# 将view张量展平为2维数组， features.size(0)获取张量的第一个维度的大小(即样本数量)
# -1 表示自动计算剩余维度的大小，这样可以确保features的形状变为[total_sample]和[feature_dim]
features_np = features.view(features.size(0), -1).numpy()
#使用 t-SNE进行降维
# 创建一个t-SNE对象，指定将数据降维到2维，n_components=2, 表示将高维数据投影到二维空间中进行可视化
# perplexity=30, 困惑度(perplexity)是t-SNE的重要参数，它影响算法对局部结构和全局结构的平衡
#                困惑度较大，t-SNE越关注全局结构，困惑度越小，t-SNE越关注局部结构，通常困惑度取值范围为5到50之间
# n_iter=300 指定n_iter算法的最大迭代次数
tsne = TSNE(n_components=2, perplexity=30, n_iter=300)
#对features_np进行t-SNE降维，返回降维后的数据tsne-sample
# tsne_results形状为[total_sample, 2]的numpy数组，表示每个样本在二维空间中的坐标
tsne_results = tsne.fit_transform(features_np)

# tsne_results是形状为[total_sample, 2]的numpy数组，表示每个样本在二维空间中的坐标
# tsne_results[:, 0] 取tsne_results的第一列，即每个样本在二维空间中的x坐标
# tsne_results[:, 1] 取tsne_results的第二列，即每个样本在二维空间中的y坐标
# c=labels 使用labels作为颜色编码，labels是包含每个样本真实标签的数组，不同的标签会有不同的颜色表示，可以通过颜色来区分样本类别
plt.figure(figsize=(10, 8))
plt.scatter(tsne_results[:, 0], tsne_results[:, 1], c=labels, cmap="viridis", s=10)
plt.colorbar()
plt.title('t_SNE Visualization after Model Training')
plt.xlabel('t-SNE Dimension 1')
plt.ylabel('t-SNE Dimension 2')
plt.show()